import numpy as np
from common.img_utils.sas_python.utils import sparse
from scipy.linalg import eig
from sklearn.cluster import KMeans, MiniBatchKMeans


def Tcut(B, Nseg, img_size, thres_evec_shape=20000, n_init=5):
    """Transfer Cut of a Bipartite Graph

    Reference:
    Li, Z., Wu, X. M., & Chang, S. F. (2012, June).
    Segmentation using superpixels: A bipartite graph partitioning approach.
    In 2012 IEEE conference on computer vision and pattern recognition (pp. 789-796). IEEE.
    https://github.com/ColumbiaDVMM/Segmentation-Using-Superpixels/blob/master/algorithms/Tcut.m

    Parameters
    ----------
    B : numpy array
        The bipartite graph
    Nseg : int
        number of segmentation
    img_size : list
        [h, w]
    thres_evec_shape : int, optional
        threshold of the amount of eigen vectors for KMeans algorithm, if less then this threshold,
        perform the full version of KMmeans, otherwise perform the problistic version of KMeans, by default 20000
    n_init : int, optional
        Number of time the k-means algorithm will be run with different centroid seeds.
        The final results will be the best output of n_init consecutive runs in terms of inertia, by default 5

    Returns
    -------
    label_img : ndarray
        the SAS segmentation label map

    Raises
    ------
    Exception
        the number of generated superpixels is too small, unable to perform the Transfer Cut
    """

    # B - |X|-by-|Y|, cross-affinity-matrix
    # note that |X| = |Y| + |I|

    Nx, Ny = B.shape
    if Ny < Nseg:
        raise Exception("Need more superpixels!")

    # Build the superpixel graph
    dx = B.sum(axis=1)
    dx_T = dx.T
    dx_T = np.array(dx_T)[0]  # convert to 1-d shape
    Nx_arr = np.array(list(range(Nx)))
    Dx = sparse(Nx_arr, Nx_arr, 1 / dx_T)
    Wy = B.T * Dx * B

    # Compute Ncut eigenvectors
    # normalized affinity matrix
    d = Wy.sum(axis=1)
    d_T = d.T
    d_T = np.array(d_T)[0]
    Ny_arr = np.array(list(range(Ny)))
    D = sparse(Ny_arr, Ny_arr, 1 / np.sqrt(d_T))
    nWy = D * Wy * D
    nWy_T = nWy.T
    nWy = (nWy + nWy_T) / 2

    # Computer eigenvectors
    print("Computer eigenvectors...")
    nWy_arr = nWy.todense()
    eval, evec = eig(nWy_arr)
    idx = np.argsort(eval)[::-1]
    Ncut_evec = D * evec[:, idx[:Nseg]]

    # Compute the Ncut eigenvectors on the entire bipartite graph (transfer!)
    evec = Dx * B * Ncut_evec

    # K-means clustering
    # extract spectral representations for pixels
    Np = np.array(img_size).prod()
    evec = evec[:Np, :]

    # Normalize each row to unit norm
    evec_pow = np.multiply(evec, evec)
    evec_divisor = np.sqrt(np.sum(evec_pow, axis=1)) + 1e-10
    evec_divisor = np.expand_dims(evec_divisor, axis=1)
    evec = evec / evec_divisor

    # K-means
    print("Process K-means...")
    if evec.shape[0] <= thres_evec_shape:
        k_means = KMeans(n_clusters=Nseg, n_init=n_init)
    else:
        k_means = MiniBatchKMeans(n_clusters=Nseg, n_init=n_init)
    evec_real = evec.real
    k_means.fit(evec_real)
    labels = k_means.labels_

    label_img = np.reshape(labels, img_size[::-1])
    label_img = label_img.T

    return label_img
